{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4c47495",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import json\n",
    "import scipy.sparse.linalg as sla\n",
    "import matplotlib.pyplot as plt\n",
    "from pathlib import Path\n",
    "from matplotlib import cm\n",
    "from tqdm import tqdm\n",
    "\n",
    "from models import *\n",
    "from modules import *\n",
    "from utils.numerical_method import *\n",
    "from utils.delta import normal\n",
    "from utils.myplot import *\n",
    "from utils.test_functions import *\n",
    "from utils.datahandler import *\n",
    "\n",
    "def _unpack_ckpt_folder(name):\n",
    "    equation_type, input_type, method, n, model_name, bs, label = name.split('_')\n",
    "    n = int(n)\n",
    "    bs = int(bs)\n",
    "    net_args = model_name.split('-')\n",
    "    net = model_names[net_args[0]](*[int(v) for v in net_args[1:]]).float()\n",
    "    return equation_type, input_type, method, n, net, bs, label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f64ff51",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Generate Data for training\n",
    "# from utils.datahandler import gen_data\n",
    "# grid_sizes = [256]\n",
    "# for n in grid_sizes:\n",
    "#     gen_data(n, train_N=2000, label='Sample10000')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "faa2eb13",
   "metadata": {},
   "outputs": [],
   "source": [
    "import gc\n",
    "import json\n",
    "\n",
    "def finite_solve_cases(equation_type, points, n):\n",
    "    F = assemble_F(n, points)\n",
    "    h = 2 / (n - 1)\n",
    "    match equation_type:\n",
    "        case 'Lap':\n",
    "            A = sparse.csr_matrix(lap_A(n))\n",
    "            B = np.stack(list(lap_b(f.squeeze(), h) for f in F), axis=1)\n",
    "\n",
    "        case 'Reaction':\n",
    "            A = sparse.csc_matrix(reaction_A(n))\n",
    "            B = np.stack(list(reaction_b(f.squeeze(), h) for f in F), axis=1)\n",
    "\n",
    "    U = sla.spsolve(A, B).transpose().reshape(len(F), n, n)\n",
    "    return U\n",
    "\n",
    "def predict_one_ckpt(cases, exp_name, js, best_or_last):\n",
    "    equation_type, data_type, method, n, net, bs, label= _unpack_ckpt_folder(exp_name)\n",
    "    ckpt_path = '/'.join(js[exp_name][\"best_model_ckpt\"].split('/')[:-1])\n",
    "    match equation_type:\n",
    "        case 'Lap':\n",
    "            pl_module = LapModule.load_from_checkpoint(\n",
    "                checkpoint_path=f'{ckpt_path}/{best_or_last}.ckpt',\n",
    "                net = net,\n",
    "                n = n)\n",
    "\n",
    "        case 'Reaction':\n",
    "            pl_module = ReactionModule.load_from_checkpoint(              \n",
    "                checkpoint_path=f'{ckpt_path}/{best_or_last}.ckpt',\n",
    "                net = net,\n",
    "                n = n)\n",
    "    pl_module.eval()\n",
    "\n",
    "    # Test all cases\n",
    "    U = finite_solve_cases(equation_type, cases, n).reshape(-1, n, n)\n",
    "    F = assemble_F(n, cases)\n",
    "    tensor = assemble_tensor(F, data_type, cases)\n",
    "    with torch.no_grad():\n",
    "        tensor = tensor.cuda().float()\n",
    "        pl_module.cuda()\n",
    "        pres = pl_module(tensor)\n",
    "        pres = pl_module.padder(pres).cpu().numpy().reshape(-1, n, n)\n",
    "    return n, F.reshape(-1, n, n), pres, U.reshape(-1, n, n)\n",
    "\n",
    "def test_all_logs_json(cases, js, names, imgs_save_path, best_or_last='best'):\n",
    "    gridSizes, Forces, U, V, Diff = [], [], [], [], []\n",
    "    for exp in names:\n",
    "        n, F, pre, ans = predict_one_ckpt(cases, exp, js, best_or_last)\n",
    "        l2_errors = cal_l2(pre, ans, n)\n",
    "\n",
    "        gridSizes.append(n)\n",
    "        Forces.append(F)\n",
    "        U.append(ans)\n",
    "        V.append(pre)\n",
    "        Diff.append(np.abs(pre - ans))\n",
    "\n",
    "        gc.collect()\n",
    "        torch.cuda.empty_cache()\n",
    "\n",
    "        latex_code = '  &  '.join([exp] + [f\"${e:.3e}$\" for e in l2_errors])\n",
    "        print(latex_code)\n",
    "\n",
    "    diff_levels = np.linspace(0.0, max(Diff), 50)\n",
    "    for i, exp in enumerate(names):\n",
    "        n = gridSizes[i]\n",
    "        save_green_images(imgs_save_path, \n",
    "            Forces[i].reshape(-1, n, n), \n",
    "            V[i].reshape(-1, n, n), \n",
    "            U[i].reshape(-1, n, n), \n",
    "            Diff[i].reshape(-1, n, n), \n",
    "            exp,\n",
    "            diff_levels)   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e059ace2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test Green function\n",
    "cases = [(0, 0), (0.5, 0), (-0.75, -0.75)]\n",
    "js_name = 'data_amount'\n",
    "\n",
    "with open(f'./js_logs/{js_name}.json') as f:\n",
    "    js = json.load(f)\n",
    "\n",
    "exp_names = [k for k in js.keys() if '2000' in k and '64' in k ]\n",
    "# exp_names = [\n",
    "#     'Lap_xyf_jac-adaptive-mul-40_64_VaryUNet-3-1-32-5_5_Sample2000'\n",
    "# ]\n",
    "\n",
    "for tag in ['best_val']:\n",
    "    print('*'*150)\n",
    "    print(tag)\n",
    "    print('*'*150)\n",
    "    test_all_logs_json(\n",
    "        cases, \n",
    "        js=js,\n",
    "        names = exp_names,\n",
    "        imgs_save_path=f'./images/{js_name}',\n",
    "        best_or_last='best_val')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba5821ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _predict_points(pl_module, n, input_type, points):\n",
    "    F = assemble_F(n, points)\n",
    "    tensor = assemble_tensor(F, input_type, points)\n",
    "    with torch.no_grad():\n",
    "        tensor = tensor.cuda().float()\n",
    "        pl_module.cuda()\n",
    "        pres = pl_module(tensor)\n",
    "        pres = pl_module.padder(pres).cpu().numpy().reshape(-1, n, n)\n",
    "    return pres \n",
    "\n",
    "def gen_green_func(exp_name, ckpt_path):\n",
    "    equation_type, input_type, method, n, net, bs, label = _unpack_ckpt_folder(exp_name)\n",
    "    match equation_type:\n",
    "        case 'Lap':\n",
    "            pl_module = LapModule.load_from_checkpoint(\n",
    "                checkpoint_path=f\"{ckpt_path}\",\n",
    "                net = net,\n",
    "                n = n)\n",
    "        case 'Reaction':\n",
    "            pl_module = ReactionModule.load_from_checkpoint(\n",
    "                checkpoint_path=f\"{ckpt_path}\",\n",
    "                net = net,\n",
    "                n = n)\n",
    "    pl_module.eval()\n",
    "    \n",
    "    def _green(points):\n",
    "        pres = _predict_points(pl_module, n, input_type, points)\n",
    "        return pres\n",
    "    return _green\n",
    "\n",
    "def sym_homo_integral(name, ckpt_path, source_func, a=1, n=64, m=32):\n",
    "    green = gen_green_func(name, ckpt_path)\n",
    "    # green = numerical_green_func(ckpt_name)\n",
    "    u = np.zeros((m, m))\n",
    "    # Compute all points source solutions for integrate\n",
    "    x = np.linspace(-a, a, n)\n",
    "    y = np.linspace(-a, a, n)\n",
    "    xx, yy = np.meshgrid(x, y)\n",
    "    f = source_func(xx, yy)\n",
    "\n",
    "    px = np.linspace(-a, a, m)\n",
    "\n",
    "    for i in range(1, m-1):\n",
    "        for j in range(1, m-1):\n",
    "            g = f * green([(px[j], px[i])]).reshape(n, n)\n",
    "            u[i, j] = simpson_2d(x, y, g)\n",
    "            # u[i, j] = weighted_linear_int_2d(x, y, 1/(n-1), g, px[j], px[i])\n",
    "            # u[i, j] = romb_2d(x, y, g)\n",
    "\n",
    "            # u[i, j] = linear_int_2d(2 / (n + 1), g) \n",
    "            # u[i, j] = integrate2d_with_interpolation(x, y, g)\n",
    "            # print(u[i, j])\n",
    "\n",
    "    return u\n",
    "\n",
    "def sym_inhomo_integral(name, ckpt_path, source_func, boundary_func, a, n, m, func_a=None):\n",
    "    green = gen_green_func(name, ckpt_path)\n",
    "\n",
    "    u = np.zeros((m, m))\n",
    "    \n",
    "    # Compute all points source solutions for integrate\n",
    "    x = np.linspace(-a, a, n)\n",
    "    y = np.linspace(-a, a, n)\n",
    "    xx, yy = np.meshgrid(x, y)\n",
    "    f = source_func(xx, yy)\n",
    "    if func_a is not None:\n",
    "        mat_a = func_a(xx, yy)\n",
    "    else:\n",
    "        mat_a = np.ones_like(f)\n",
    "    px = np.linspace(-a, a, m)\n",
    "    v = np.linspace(-a, a, n)\n",
    "    h = 2 * a / (n - 1)\n",
    "    for i in range(1, m-1):\n",
    "        for j in range(1, m-1):\n",
    "            g = green([(px[j], px[i])]).reshape(n, n)\n",
    "            gf = np.float128(f * g)\n",
    "            uf = simpson_2d(x, y, gf) \n",
    "            # uf = linear_int_2d(2 / (n + 1), gf) \n",
    "\n",
    "            # top side\n",
    "            ug_top = boundary_func(x=v, y=a)\n",
    "            va_top = mat_a[-1, :]\n",
    "\n",
    "            flow_top = (g[-1, :] - g[-2, :]) /  h\n",
    "            # flow_top = (3*g[-1, :] - 4*g[-2, :] + g[-3, :]) / (2 * h)\n",
    "            \n",
    "            # val_top = romb((va_top * ug_top * flow_top)[:-1], h)\n",
    "            val_top = simps((va_top * ug_top * flow_top)[:-1], v[:-1])\n",
    "            # val_top = (va * ug_top * flow_top * h).sum()\n",
    "            # val_top = weighted_linear_int_1d(v, h, va_top * ug_top * flow_top, px[j])\n",
    "            \n",
    "            # Bottom side\n",
    "            ug_bottom = boundary_func(x=v, y=-a)\n",
    "            va_bottom = mat_a[0, :]\n",
    "            \n",
    "            flow_bottom = (g[0, :] - g[1, :]) /  h\n",
    "            # flow_bottom = (3*g[0, :] - 4*g[1, :] + g[2, :]) / (2 * h)\n",
    "            \n",
    "            # val_bottom = romb((va_bottom * ug_bottom * flow_bottom)[1:], h)\n",
    "            val_bottom = simps((va_bottom * ug_bottom * flow_bottom)[1:], v[1:])\n",
    "            # val_bottom = (va * ug_bottom * flow_bottom * h).sum()\n",
    "            # val_bottom = weighted_linear_int_1d(v, h, va_bottom * ug_bottom * flow_bottom, px[j])\n",
    "\n",
    "            # left side\n",
    "            ug_left = boundary_func(x=-a, y=v)\n",
    "            va_left = mat_a[:, 0]\n",
    "            \n",
    "            flow_left = (g[:, 0] - g[:, 1]) / h\n",
    "            # flow_left = (3*g[:, 0] - 4*g[:, 1] + g[:, 2]) / (2 * h)\n",
    "            \n",
    "            # val_left = romb((va_left * ug_left * flow_left)[:-1], h)\n",
    "            val_left = simps((va_left * ug_left * flow_left)[:-1], v[:-1])\n",
    "            # val_left = (va * ug_left * flow_left * h).sum()\n",
    "            # val_left = weighted_linear_int_1d(v, h, va_left * ug_left * flow_left, px[i])\n",
    "\n",
    "            # right side\n",
    "            ug_right = boundary_func(x=a, y=v)\n",
    "            va_right = mat_a[:, -1]\n",
    "            \n",
    "            flow_right = (g[:, -1] - g[:, -2]) / h\n",
    "            # flow_right = (3*g[:, -1] - 4*g[:, -2] + g[:, -3]) / (2 * h)\n",
    "            \n",
    "            # val_right = romb((va_right * ug_right * flow_right)[1:], h)\n",
    "            val_right = simps((va_right * ug_right * flow_right)[1:], v[1:])\n",
    "            # val_right = (va * ug_right * flow_right * h).sum()\n",
    "            # val_right = weighted_linear_int_1d(v, h, va_right * ug_right * flow_right, px[i])\n",
    "\n",
    "\n",
    "            u[i, j] = uf - val_top - val_bottom - val_left - val_right\n",
    "            # u[i, j] = uf\n",
    "\n",
    "    return u\n",
    "\n",
    "def integral_save_img(name, ckpt_path, img_save_path, func_force, func_ref, func_boundary=None, label='f1', sample_nums=[64, 32]):\n",
    "    equation_type, input_type, method, n, net, bs, _ = _unpack_ckpt_folder(name)\n",
    "    \n",
    "    def _save(name, n, func_force, func_ref, pre):\n",
    "        x = np.linspace(-1, 1, n)[1:-1]\n",
    "        y = np.linspace(-1, 1, n)[1:-1]\n",
    "        xx, yy = np.meshgrid(x, y)\n",
    "        f = func_force(xx, yy)\n",
    "        ans = func_ref(xx, yy)\n",
    "\n",
    "        h = 2 / (n-1)\n",
    "        l2 = func_L2norm(pre, ans, h)\n",
    "        save_FastSolver_images(Path(f'{name}_{l2:.3e}'), f, pre, ans)\n",
    "        return l2\n",
    "    \n",
    "    for sample_num in sample_nums:\n",
    "        if func_boundary is None:\n",
    "            pre = sym_homo_integral(name, ckpt_path, func_force, 1, n, sample_num)\n",
    "            l2 = _save(f'{img_save_path}/sym_{label}_{sample_num}', sample_num, func_force, func_ref, pre[1:-1, 1:-1])\n",
    "        else:\n",
    "            if equation_type == 'Reaction':\n",
    "                function_a = func_a\n",
    "            else:\n",
    "                function_a = None\n",
    "            pre = sym_inhomo_integral(name, ckpt_path, func_force, func_boundary, 1, n, sample_num, function_a,)\n",
    "            l2 = _save(f'{img_save_path}/sym_{label}_{sample_num}', sample_num, func_force, func_ref, pre[1:-1, 1:-1], )\n",
    "    return f\"{label} & {l2:.3e}\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d44b79a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "funcs = [f0, f1, f2, f3, f4, f5]\n",
    "sols = [u0, u1, u2, u3, u4, u5]\n",
    "\n",
    "js_name = 'data_amount'\n",
    "\n",
    "with open(f'./js_logs/{js_name}.json') as f:\n",
    "    js = json.load(f)\n",
    "\n",
    "exp_names = [k for k in js.keys() if 'Lap' in k and '2000' in k and '128' in k]\n",
    "sample_nums = [64]\n",
    "for name in exp_names:\n",
    "    if 'Lap' in name:\n",
    "        ckpt_path = '/'.join(js[name][\"best_model_ckpt\"].split('/')[:-1])\n",
    "        print(name)\n",
    "        for tag in ['best_val']:\n",
    "            print(\"*\"*100)\n",
    "            print(tag)\n",
    "            print(\"*\"*100)\n",
    "\n",
    "            p = f\"{ckpt_path}/{tag}.ckpt\"\n",
    "            for i, (f, u) in enumerate(zip(funcs, sols)):\n",
    "                if i < 2:\n",
    "                    latex_info = integral_save_img(name, p, f'images/{js_name}/{name}', f, u, None, f'f{i+1}', sample_nums)\n",
    "                else:\n",
    "                    latex_info = integral_save_img(name, p, f'images/{js_name}/{name}', f, u, u, f'f{i+1}', sample_nums)\n",
    "                print(latex_info)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c1a1dab",
   "metadata": {},
   "outputs": [],
   "source": [
    "js_name = 'review_reaction'\n",
    "\n",
    "with open(f'./js_logs/{js_name}.json') as f:\n",
    "    js = json.load(f)\n",
    "\n",
    "exp_names = [k for k in js.keys() if 'Reaction' in k]\n",
    "sample_nums = [64]\n",
    "\n",
    "for name in exp_names:\n",
    "    ckpt_path = '/'.join(js[name][\"best_model_ckpt\"].split('/')[:-1])\n",
    "    print(\"\\n\\n\", name)\n",
    "    for tag in ['best_val']:\n",
    "        print(\"*\"*100)\n",
    "        print(tag)\n",
    "        print(\"*\"*100)\n",
    "\n",
    "        p = f\"{ckpt_path}/{tag}.ckpt\"\n",
    "        latex_info =integral_save_img(name, p, f'images/{js_name}/{name}', force, func_u, func_u, 'inhomo', sample_nums)\n",
    "        print(latex_info)\n",
    "\n",
    "        latex_info =integral_save_img(name, p, f'images/{js_name}/{name}',f6, u6, None, 'homo', sample_nums)\n",
    "        print(latex_info)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1165786",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "csv_path = Path('./csv')\n",
    "logs = []\n",
    "for i, num in enumerate([1000, 2000, 5000, 10000]):\n",
    "    f = csv_path / f\"Lap_xyf_jac-const-5_64_VaryUNet-3-1-16-4_5_{num}_version_0.csv\"\n",
    "    df = pd.read_csv(f, usecols=['Value'])\n",
    "    df = df.rename({'Value': r\"$S_{%d}$\" %(i+1)}, axis=1)\n",
    "    logs.append(df)\n",
    "plt.rcParams[\"font.size\"] = 24\n",
    "df = pd.concat(logs, axis=1)\n",
    "df = df.rolling(60, step=60).mean()\n",
    "fig = df.plot(\n",
    "    figsize=(23, 9),\n",
    "    title='Validation',\n",
    "    logy=True,\n",
    "    style=['b-*', 'g-o', 'r-x', '-|'],\n",
    "    xlabel='Epoch',\n",
    "    ylabel='Error'\n",
    ")\n",
    "plt.grid(True)\n",
    "plt.xticks([100 * i for i in range(11)], [r\"%d\" %(15 * i) for i in range(11)])\n",
    "\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  },
  "vscode": {
   "interpreter": {
    "hash": "51bec75535e8b8e121f5058573fd16f5d4f5054d3fb39f60867bdc80b6d322e7"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
